#!/bin/bash
# Coze Studio 一键安装与更新脚本
# 适用于 CentOS 7 系统
# 版本: 1.0

# 配置参数
REPO_URL="https://gitee.com/wmob/coze-studio-1.git"
PROJECT_DIR="/opt/coze-studio"
VENV_DIR="${PROJECT_DIR}/venv"
DB_NAME="coze_db"
DB_USER="coze_user"
DB_PASSWORD=$(openssl rand -hex 12)  # 随机生成数据库密码
SECRET_KEY=$(openssl rand -hex 32)  # 随机生成应用密钥
PYTHON_VERSION="3.9.17"
PYTHON_TAR="Python-${PYTHON_VERSION}.tgz"
PYTHON_URL="https://www.python.org/ftp/python/${PYTHON_VERSION}/${PYTHON_TAR}"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # 无颜色

# 检查是否以 root 身份运行
if [ "$(id -u)" -ne 0 ]; then
    echo -e "${RED}错误: 请使用 root 权限运行此脚本${NC}" >&2
    exit 1
fi

# 显示菜单
show_menu() {
    echo -e "\n${YELLOW}Coze Studio 安装管理工具${NC}"
    echo "1. 全新安装"
    echo "2. 更新系统"
    echo "3. 启动服务"
    echo "4. 停止服务"
    echo "5. 重启服务"
    echo "6. 查看状态"
    echo "7. 退出"
    echo -n "请选择操作 [1-7]: "
}

# 安装基础依赖
install_dependencies() {
    echo -e "\n${YELLOW}正在安装系统依赖...${NC}"
    yum update -y
    yum install -y git wget gcc gcc-c++ make openssl-devel bzip2-devel libffi-devel \
                   postgresql13 postgresql13-server redis nginx firewalld
}

# 安装 Python 3.9
install_python() {
    echo -e "\n${YELLOW}正在安装 Python ${PYTHON_VERSION}...${NC}"
    if ! command -v python3.9 &> /dev/null; then
        wget -q ${PYTHON_URL} -O ${PYTHON_TAR}
        tar -zxf ${PYTHON_TAR}
        cd Python-${PYTHON_VERSION}
        ./configure --prefix=/usr/local/python39 --with-openssl=/usr/include/openssl
        make -j $(nproc)
        make install
        ln -s /usr/local/python39/bin/python3.9 /usr/bin/python3
        ln -s /usr/local/python39/bin/pip3.9 /usr/bin/pip3
        cd ..
        rm -rf Python-${PYTHON_VERSION} ${PYTHON_TAR}
    else
        echo -e "${GREEN}Python ${PYTHON_VERSION} 已安装${NC}"
    fi
}

# 配置数据库
configure_database() {
    echo -e "\n${YELLOW}正在配置数据库...${NC}"
    # 初始化 PostgreSQL
    if [ ! -d "/var/lib/pgsql/13/data" ]; then
        /usr/pgsql-13/bin/postgresql-13-setup initdb
    fi
    
    # 启动并设置开机自启
    systemctl enable --now postgresql-13
    
    # 创建数据库和用户
    sudo -u postgres psql -c "CREATE DATABASE IF NOT EXISTS ${DB_NAME};"
    sudo -u postgres psql -c "CREATE USER IF NOT EXISTS ${DB_USER} WITH PASSWORD '${DB_PASSWORD}';"
    sudo -u postgres psql -c "GRANT ALL PRIVILEGES ON DATABASE ${DB_NAME} TO ${DB_USER};"
}

# 配置 Redis
configure_redis() {
    echo -e "\n${YELLOW}正在配置 Redis...${NC}"
    systemctl enable --now redis
}

# 配置 Nginx
configure_nginx() {
    echo -e "\n${YELLOW}正在配置 Nginx...${NC}"
    cat > /etc/nginx/conf.d/coze.conf << EOF
server {
    listen 80;
    server_name _;

    location / {
        proxy_pass http://127.0.0.1:8000;
        proxy_set_header Host \$host;
        proxy_set_header X-Real-IP \$remote_addr;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto \$scheme;
    }
}
EOF
    systemctl enable --now nginx
    systemctl restart nginx
}

# 配置防火墙
configure_firewall() {
    echo -e "\n${YELLOW}正在配置防火墙...${NC}"
    systemctl start firewalld
    systemctl enable firewalld
    firewall-cmd --zone=public --add-port=80/tcp --permanent
    firewall-cmd --reload
}

# 克隆/更新代码仓库
clone_or_update_repo() {
    echo -e "\n${YELLOW}正在获取项目代码...${NC}"
    if [ -d "${PROJECT_DIR}/.git" ]; then
        echo -e "${YELLOW}更新现有代码...${NC}"
        cd ${PROJECT_DIR}
        git pull
    else
        echo -e "${YELLOW}克隆新代码...${NC}"
        git clone ${REPO_URL} ${PROJECT_DIR}
    fi
}

# 配置项目环境
configure_project() {
    echo -e "\n${YELLOW}正在配置项目环境...${NC}"
    # 创建虚拟环境
    if [ ! -d "${VENV_DIR}" ]; then
        python3 -m venv ${VENV_DIR}
    fi
    
    # 激活虚拟环境并安装依赖
    source ${VENV_DIR}/bin/activate
    pip3 install --upgrade pip
    pip3 install -r ${PROJECT_DIR}/requirements.txt
    pip3 install gunicorn
    
    # 创建环境变量配置文件
    cat > ${PROJECT_DIR}/.env << EOF
# 数据库配置
DB_HOST=localhost
DB_PORT=5432
DB_NAME=${DB_NAME}
DB_USER=${DB_USER}
DB_PASSWORD=${DB_PASSWORD}

# Redis 配置
REDIS_HOST=localhost
REDIS_PORT=6379
REDIS_DB=0

# 应用配置
SECRET_KEY=${SECRET_KEY}
DEBUG=False
ALLOWED_HOSTS=*
PORT=8000
EOF
    
    # 数据库迁移
    cd ${PROJECT_DIR}
    if [ -f "manage.py" ]; then
        python3 manage.py migrate
        # 检查超级用户是否存在，不存在则创建
        if ! python3 manage.py shell -c "from django.contrib.auth import get_user_model; User = get_user_model(); print(User.objects.filter(is_superuser=True).exists())" | grep -q "True"; then
            echo -e "${YELLOW}请创建管理员账号:${NC}"
            python3 manage.py createsuperuser
        fi
    else
        echo -e "${YELLOW}未找到 manage.py，尝试 Flask 迁移...${NC}"
        if [ -f "app.py" ] || [ -f "wsgi.py" ]; then
            flask db init 2>/dev/null
            flask db migrate -m "initial migration"
            flask db upgrade
        fi
    fi
    
    deactivate
}

# 配置系统服务
configure_systemd() {
    echo -e "\n${YELLOW}正在配置系统服务...${NC}"
    cat > /etc/systemd/system/coze.service << EOF
[Unit]
Description=Coze Studio Service
After=network.target postgresql-13.service redis.service nginx.service

[Service]
User=root
Group=root
WorkingDirectory=${PROJECT_DIR}
ExecStart=${VENV_DIR}/bin/gunicorn --workers 4 --bind 127.0.0.1:8000 wsgi:application
Restart=always
RestartSec=5
Environment="PATH=${VENV_DIR}/bin"

[Install]
WantedBy=multi-user.target
EOF
    
    systemctl daemon-reload
    systemctl enable coze
}

# 全新安装
full_install() {
    echo -e "${YELLOW}开始全新安装 Coze Studio...${NC}"
    install_dependencies
    install_python
    configure_database
    configure_redis
    configure_nginx
    configure_firewall
    clone_or_update_repo
    configure_project
    configure_systemd
    start_service
    echo -e "\n${GREEN}Coze Studio 安装完成!${NC}"
    echo -e "${GREEN}您可以通过服务器 IP 地址访问系统${NC}"
}

# 更新系统
update_system() {
    echo -e "${YELLOW}开始更新 Coze Studio...${NC}"
    stop_service
    install_dependencies
    clone_or_update_repo
    source ${VENV_DIR}/bin/activate
    pip3 install --upgrade pip
    pip3 install -r ${PROJECT_DIR}/requirements.txt
    cd ${PROJECT_DIR}
    if [ -f "manage.py" ]; then
        python3 manage.py migrate
    else
        flask db upgrade
    fi
    deactivate
    start_service
    echo -e "${GREEN}Coze Studio 更新完成!${NC}"
}

# 启动服务
start_service() {
    echo -e "${YELLOW}启动 Coze Studio 服务...${NC}"
    systemctl start coze
    check_status
}

# 停止服务
stop_service() {
    echo -e "${YELLOW}停止 Coze Studio 服务...${NC}"
    systemctl stop coze
    check_status
}

# 重启服务
restart_service() {
    echo -e "${YELLOW}重启 Coze Studio 服务...${NC}"
    systemctl restart coze
    check_status
}

# 查看状态
check_status() {
    echo -e "${YELLOW}Coze Studio 服务状态:${NC}"
    systemctl status coze --no-pager
}

# 主程序
while true; do
    show_menu
    read -r choice
    case $choice in
        1)
            full_install
            ;;
        2)
            update_system
            ;;
        3)
            start_service
            ;;
        4)
            stop_service
            ;;
        5)
            restart_service
            ;;
        6)
            check_status
            ;;
        7)
            echo -e "${GREEN}退出程序${NC}"
            exit 0
            ;;
        *)
            echo -e "${RED}无效选择，请输入 1-7 之间的数字${NC}"
            ;;
    esac
    echo -e "\n${YELLOW}按任意键继续...${NC}"
    read -r -n 1
done
